import socket
import asyncio
from typing import Callable
import asyncio
from pathlib import Path
import sys
BASE_DIR = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(BASE_DIR).replace("\\", "\\\\"))
from agent.pkt import Pkt, Cmd
from agent.session import Session
from agent.session_admin import SessionAdmin
from agent.kcp_py.kcp import KCPClientAsync

DataHandler = Callable[[bytes], None]
import sys

sa = SessionAdmin()

class AgentClient(object):
  global_id = 0
  read_buf_max = 2048
  send_cb:DataHandler = None

  def __init__(self, local_host = '127.0.0.1', local_port = 8080, send_cb :DataHandler = None):
    self.local_host = local_host
    self.local_port = local_port
    AgentClient.send_cb = send_cb

  async def do_keepalive(self):
    while True:
      print('send keepalive')
      pkt = Pkt(AgentClient.global_id, 0, Cmd.KEEP_ALIVE, None)
      AgentClient.send_cb(pkt.pack())
      await asyncio.sleep(10)

  async def on_start(self):
    print('listen local ', self.local_host, self.local_port)
    print("Connected to server for global id!")
    while True:
      if AgentClient.global_id != 0:
        break
      pkt = Pkt(0, 0, Cmd.QUERY_GLOBALID)
      AgentClient.send_cb(pkt.pack())
      await asyncio.sleep(1)

    loop = asyncio.get_running_loop()
    loop.create_task(self.do_keepalive())
    await asyncio.start_server(self.handle_cb, self.local_host, self.local_port)

  async def on_data(self, data :bytes):
    pkt = Pkt(0, 0, 0)
    if pkt.unpack(data) is False:
      print('recv err data')
      return
    if pkt.cmd == Cmd.QUERY_GLOBALID:
      AgentClient.global_id = pkt.globalid
      print('get global_id', AgentClient.global_id)
      return
    if pkt.cmd == Cmd.KEEP_ALIVE:
      print('get keepalive')
      return
    s = sa.get_session_by_ids(pkt.globalid, pkt.localid)
    if s is None:
      print('not found session by ids ', pkt.globalid, pkt.localid)
      return
    
    s.writer.write(pkt.data)

  @staticmethod
  def get_global_id():
    return AgentClient.global_id

  @staticmethod    
  def handle_cb(reader, writer):
    print('agent recv !!!!')
    s = Session(AgentClient.get_global_id(), sa.get_add_session_id(), reader, writer)
    sa.add_session(s)
    loop = asyncio.get_running_loop()
    loop.create_task(AgentClient.client_proxy_proc(s))

  @staticmethod
  async def client_proxy_proc(s :Session):
    while True:
      try:
        data = await s.reader.read(AgentClient.read_buf_max)
        if data is None or len(data) <= 0:
          break
        pkt = Pkt(s.global_id, s.session_id, Cmd.COMMON_DATA, data)
        AgentClient.send_cb(pkt.pack())
      except BaseException as e:
        print(e)
        break
    await asyncio.sleep(5)
    sa.del_session(s)
    

client = KCPClientAsync(
    "liuchenxi.love",
    # "127.0.0.1",
    5100,
    conv_id=1,
)

def send_cb(data :bytes) -> None:
  try:
    client.send(data)
  except BaseException as e:
    print(e)

agent = AgentClient(
  "0.0.0.0",
  8080,
  send_cb,
)

@client.on_data
async def handle_data(data: bytes) -> None:
  await agent.on_data(data)

@client.on_start
async def on_start() -> None:
  await agent.on_start()

client.start()




















